This pull request introduces head chunking over the QK attention dot product dimension in Multi-Head Latent Attention (MLA). By using jax.lax.scan to compute scores in blocks of attention heads, it prevents the full materialization of the query-key score tensor, substantially reducing the high-bandwidth memory (HBM) footprint for large context windows.

🔍 General Feedback

• Excellent Memory Efficiency: The introduction of chunking via sequential scan is a highly effective, mathematically correct design choice to mitigate OOMs at scale.
• Accidental Config Disabling: Please revert the change to enable_checkpointing in base.yml. Leaving checkpointing disabled by default is a major hazard for all production workloads.
• Add Parity and Unit Tests: Although manual verification is mentioned, it is highly recommended to include automated regression coverage in tests/unit/deepseek32_vs_reference_test.py. For example, a test named test_indexer_chunked_match could set qk_head_chunk_size to a divisor of the indexer heads (e.g., 16) and assert that the outputs are identical to the unchunked implementation.
• Silent Fallback: As noted in the inline review, consider raising an explicit ValueError when qk_head_chunk_size is misconfigured (e.g., if it exceeds the number of heads or does not divide them evenly), instead of falling back silently to unchunked materialization.

This pull request implements QK attention head chunking in MLA (Multi-Head Latent Attention) using jax.lax.scan to significantly reduce the HBM footprint for long contexts. While the overall design and mathematical approach are excellent and provide great memory efficiency, a critical refactoring bug was introduced where the dynamic head importance weights became undefined, which completely breaks the forward pass.

🔍 General Feedback

• Crucial Definition Restored: Ensure the weights tensor is properly calculated using the weights_proj projection layer before applying chunking to avoid a runtime NameError.
• Performance Optimizations: The proposed code changes optimize the scan loop bodies by formulating the einsum contractions directly on the chunked layouts, removing 4 transpose operations inside sequential loops.
• Testing & Verification: Please add automated unit tests (e.g., in tests/unit/deepseek32_vs_reference_test.py) that explicitly run with non-zero qk_head_chunk_size values and verify numerical parity with the unchunked version.

This pull request implements a highly effective chunked attention dot product for DeepSeek-V3.2 MLA. By chunking along the attention heads dimension during QK product evaluation, the implementation reduces the maximum HBM footprint of the intermediate attention scores tensor, allowing the model to handle much larger context lengths with native JAX primitives.

🔍 General Feedback

• High Quality and Parity: The implementation contains robust unit tests comparing the outputs of both the indexer and MLA module against a reference PyTorch implementation.
• Optimal Memory Trade-offs: The addition of qk_head_chunk_size allows users to scale and control memory-compute trade-offs cleanly.
• Clean Fallback Paths: Keeping the native unchunked branch as the default ensures backwards compatibility and zero performance regression when chunking is not enabled.

This PR introduces QK attention head chunking to MLA to reduce the HBM footprint, particularly beneficial for long context lengths. By iteratively evaluating head groups via jax.lax.scan and calculating softmax partials sequentially, the design achieves identical mathematical parity to full tensor materialization with significantly reduced memory footprint.

🔍 General Feedback

• High-Quality Design: The mathematical implementation using jax.lax.scan is precise and exactly matches the un-chunked fallback branch.
• Testing Coverage Gap: Although the PR description mentions manually verifying parity, the existing unit test suite in tests/unit/deepseek32_vs_reference_test.py does not parameterize over qk_head_chunk_size nor test any non-zero chunk sizes in CI. To ensure this features does not regress, we suggest adding qk_head_chunk_size as a test parameter in deepseek32_vs_reference_test.py as detailed below.

Suggested Test Improvements:

1. In tests/unit/deepseek32_vs_reference_test.py, update get_cfg_and_mesh to support qk_head_chunk_size:

   def get_cfg_and_mesh(config, run_name, dtype, batch_size, seq_len, attention, indexer_topk, qk_head_chunk_size=0):
     """Returns MaxText configuration and mesh."""
     cfg = pyconfig.initialize(
         [None, get_test_config_path()],
         run_name=run_name,
         enable_checkpointing=False,
         model_name="default",
         dtype=dtype,
         weight_dtype="float32",
         matmul_precision="highest",
         float32_qk_product=True,
         float32_logits=True,
         per_device_batch_size=batch_size,
         max_target_length=seq_len,
         max_prefill_predict_length=seq_len,
         attention=attention,
         indexer_topk=indexer_topk,
         qk_head_chunk_size=qk_head_chunk_size,
         **asdict(config),
     )
     # ... rest of method ...

2. Parameterize test_indexer_match to verify chunking configurations (e.g., 0, 2, 4):

   class DeepseekV32IndexerTest(DeepseekTestBase):
     """Tests for the Sparse Indexer (Top-K Selection)."""
   
     @parameterized.parameters(0, 2, 4)
     def test_indexer_match(self, qk_head_chunk_size, seq_len=8):
       """Verifies Indexer output matches PyTorch output."""
       # ...
       cfg, mesh = get_cfg_and_mesh(
           config=self.config,
           run_name="deepseek_indexer_test",
           dtype=self.dtype,
           batch_size=self.batch_size,
           seq_len=self.seq_len,
           attention="dot_product",
           indexer_topk=4,
           qk_head_chunk_size=qk_head_chunk_size,
       )
       # ...

3. Update test_mla_parity parameter list and test cases to verify different chunk sizes:

     def test_mla_parity(self, attention, seq_len, indexer_topk, check_norm=False, qk_head_chunk_size=0):
       # ...
       cfg, mesh = get_cfg_and_mesh(
           config=self.config,
           run_name="deepseek_mla_test",
           dtype=self.dtype,
           batch_size=self.batch_size,
           seq_len=self.seq_len,
           attention=attention,
           indexer_topk=indexer_topk,
           qk_head_chunk_size=qk_head_chunk_size,
       )
       # ...

This pull request implements a highly effective memory optimization for Multi-Layer Attention (MLA) in DeepSeek models by introducing QK head chunking. By utilizing jax.lax.scan to sequentially compute the query-key dot product and its softmax over head chunks, the attention scores tensor footprint is drastically reduced from $[batch, count_of_heads, q_len, kv_len]$ to $[batch, head_chunk_size, q_len, kv_len]$, allowing MaxText to scale to significantly larger context lengths without running into HBM constraints.

🔍 General Feedback

• High Code Quality: The JAX implementation of chunking over the heads dimension is exceptionally clean, well-documented, and mathematically identical to the full tensor materialization.
• Rematerialization Strategy: The use of jax.checkpoint on the scan body of the sparse indexer is extremely clever and essential to prevent the compiler from caching the intermediate logits, thereby successfully securing the memory reductions during training backpropagation.
• Comprehensive Verification: The addition of parameterized reference tests in tests/unit/deepseek32_vs_reference_test.py validates equivalence across multiple chunk sizes (0, 2, and 4) against PyTorch reference results, ensuring high reliability.

This pull request implements head-wise attention chunking using jax.lax.scan for both the DeepSeek-V3.2 MLA sparse indexer scores computation and its KL divergence loss function. The implementation is highly efficient, mathematically equivalent to the native (un-chunked) implementation, and successfully reduces HBM footprint to avoid out-of-memory (OOM) failures under long context lengths.

🔍 General Feedback

• Excellent JAX Idiomatic Design: Leverages jax.lax.scan to sequentially process slices of head dimension, resulting in a compiled loop (rather than Python unrolling) which keeps compile times and the XLA HLO graph extremely small.
• Optimal Gradient Isolation and Checkpointing: Strategically applies jax.checkpoint on scan_body_indexer to avoid saving large intermediate logits on the backward pass, while correctly omitting checkpointing on scan_body_heads where gradients are already stopped.
• Robust Validation Guardrails: Comprehensive boundary checks (divisibility and size) are integrated on both SparseIndexer and calculate_indexer_loss blocks to prevent incorrect reshaping or compilation-level failures.
• Thorough Test Coverage: New test cases successfully verify parity under multiple sweep sizes (qk_head_chunk_size of 0, 2, 4) against the PyTorch reference implementation.